"""
书法风格分类器
提供自动识别书法风格和名家风格的功能
"""
import os
import json
import numpy as np
import logging
import pickle
from scipy.spatial.distance import cosine
from collections import defaultdict

# 导入风格模型
from style_models import (
    get_style_model, 
    STYLE_REGULAR, 
    STYLE_RUNNING,
    STYLE_CURSIVE,
    STYLE_CLERICAL,
    STYLE_SEAL,
    SUPPORTED_STYLES,
    STYLE_NAMES
)

# 配置日志
logger = logging.getLogger(__name__)

class StyleClassifier:
    """书法风格分类器"""
    
    def __init__(self, model_path=None):
        """
        初始化风格分类器
        model_path: 风格分类器模型路径，如果提供则加载已有模型
        """
        # 加载所有风格的模型
        self.style_models = {}
        for style in SUPPORTED_STYLES:
            self.style_models[style] = get_style_model(style)
        
        # 加载名家风格特征库
        self.master_features = self._load_master_features()
        
        logger.info(f"风格分类器初始化完成，加载了{len(self.style_models)}种风格模型")
    
    def _load_master_features(self):
        """加载名家风格特征库"""
        # 名家特征库路径
        master_features_path = os.path.join(os.path.dirname(__file__), 'data', 'master_features.json')
        
        # 如果文件不存在，返回空字典
        if not os.path.exists(master_features_path):
            logger.warning(f"名家特征库文件不存在: {master_features_path}")
            return {}
        
        try:
            with open(master_features_path, 'r', encoding='utf-8') as f:
                master_features = json.load(f)
            logger.info(f"成功加载名家特征库，包含{len(master_features)}位名家的特征")
            return master_features
        except Exception as e:
            logger.error(f"加载名家特征库失败: {e}")
            return {}
    
    def classify_style(self, features):
        """
        根据提取的特征判断书法风格
        返回最可能的风格及概率分布
        """
        style_scores = {}
        
        # 核心风格特征
        key_features = [
            'angle_distribution',    # 角度分布
            'curvature',             # 曲率
            'line_width_variance',   # 线宽变化
            'fourier_descriptors',   # 傅里叶描述符
            'structural_complexity'  # 结构复杂度
        ]
        
        # 计算每种风格的匹配度
        for style_name, style_model in self.style_models.items():
            # 计算特征相似度
            similarities = []
            for feature_name in key_features:
                if feature_name not in features or feature_name not in style_model.ideal_features:
                    continue
                    
                actual = features[feature_name]
                ideal = style_model.ideal_features[feature_name]
                
                # 确保可以比较
                if isinstance(actual, list) and isinstance(ideal, list):
                    # 长度可能不同，取最小长度
                    min_len = min(len(actual), len(ideal))
                    if min_len > 0:
                        actual = actual[:min_len]
                        ideal = ideal[:min_len]
                        sim = 1.0 - self._calculate_similarity(actual, ideal)
                        similarities.append(sim)
                elif not isinstance(actual, list) and not isinstance(ideal, list):
                    # 单值特征
                    diff = abs(actual - ideal) / max(abs(ideal), 1e-10)
                    sim = max(0, 1.0 - diff)
                    similarities.append(sim)
            
            if not similarities:
                style_scores[style_name] = 0.0
                continue
                
            # 计算平均匹配度
            avg_similarity = sum(similarities) / len(similarities)
            style_scores[style_name] = avg_similarity
        
        # 将分数归一化为概率分布
        total_score = sum(style_scores.values())
        if total_score > 0:
            for style in style_scores:
                style_scores[style] = style_scores[style] / total_score
        
        # 找出最匹配的风格
        best_style = max(style_scores.items(), key=lambda x: x[1])
        
        result = {
            'best_style': best_style[0],
            'best_style_name': STYLE_NAMES.get(best_style[0], best_style[0]),
            'confidence': best_style[1],
            'all_styles': {k: {'name': STYLE_NAMES.get(k, k), 'score': v} for k, v in style_scores.items()}
        }
        
        logger.info(f"风格分类结果: {result['best_style_name']}，置信度: {result['confidence']:.2f}")
        return result
    
    def find_similar_masters(self, features, top_n=3):
        """
        根据提取的特征寻找最相似的名家风格
        返回top_n个最相似的名家及相似度
        """
        if not self.master_features:
            logger.warning("名家特征库为空，无法进行名家风格匹配")
            return []
        
        master_scores = {}
        
        # 核心风格特征
        key_features = [
            'angle_distribution',    # 角度分布
            'curvature',             # 曲率
            'line_width_variance',   # 线宽变化
            'fourier_descriptors',   # 傅里叶描述符
            'structural_complexity'  # 结构复杂度
        ]
        
        # 计算与每位名家的相似度
        for master_name, master_feature_set in self.master_features.items():
            similarities = []
            
            for feature_name in key_features:
                if feature_name not in features or feature_name not in master_feature_set:
                    continue
                    
                actual = features[feature_name]
                master_feature = master_feature_set[feature_name]
                
                # 确保可以比较
                if isinstance(actual, list) and isinstance(master_feature, list):
                    # 长度可能不同，取最小长度
                    min_len = min(len(actual), len(master_feature))
                    if min_len > 0:
                        actual = actual[:min_len]
                        master_feature = master_feature[:min_len]
                        sim = 1.0 - self._calculate_similarity(actual, master_feature)
                        similarities.append(sim)
                elif not isinstance(actual, list) and not isinstance(master_feature, list):
                    # 单值特征
                    diff = abs(actual - master_feature) / max(abs(master_feature), 1e-10)
                    sim = max(0, 1.0 - diff)
                    similarities.append(sim)
            
            if not similarities:
                continue
                
            # 计算平均相似度
            avg_similarity = sum(similarities) / len(similarities)
            master_scores[master_name] = avg_similarity
        
        # 按相似度降序排序
        sorted_masters = sorted(master_scores.items(), key=lambda x: x[1], reverse=True)
        
        # 返回top_n结果
        top_masters = [
            {'name': name, 'similarity': score} for name, score in sorted_masters[:top_n]
        ]
        
        if top_masters:
            logger.info(f"最相似的名家: {top_masters[0]['name']}，相似度: {top_masters[0]['similarity']:.2f}")
        
        return top_masters
    
    def _calculate_similarity(self, v1, v2):
        """计算两个向量的相似度，返回0-1之间的值，越小越相似"""
        # 确保向量长度一致
        min_len = min(len(v1), len(v2))
        v1 = v1[:min_len]
        v2 = v2[:min_len]
        
        # 转换为numpy数组
        v1 = np.array(v1)
        v2 = np.array(v2)
        
        # 如果向量全为零，返回最大距离
        if np.sum(np.abs(v1)) < 1e-10 or np.sum(np.abs(v2)) < 1e-10:
            return 1.0
        
        try:
            # 使用余弦距离
            similarity = cosine(v1, v2)
            # 确保结果在0-1之间
            return min(max(0.0, similarity), 1.0)
        except:
            # 如果出错，返回最大距离
            return 1.0
    
    def add_master_example(self, master_name, features, style=None):
        """
        添加名家作品示例到特征库
        master_name: 名家名称
        features: 特征向量
        style: 书法风格
        """
        if not self.master_features:
            self.master_features = {}
        
        # 如果名家不存在，创建新条目
        if master_name not in self.master_features:
            self.master_features[master_name] = {
                'style': style,
                'features': []
            }
        
        # 添加特征
        self.master_features[master_name]['features'].append(features)
        
        # 更新风格信息
        if style:
            self.master_features[master_name]['style'] = style
        
        logger.info(f"已添加{master_name}的作品示例到特征库")
        
        # 保存特征库
        self._save_master_features()
        
        return True
    
    def _save_master_features(self):
        """保存名家特征库"""
        # 名家特征库路径
        master_features_path = os.path.join(os.path.dirname(__file__), 'data', 'master_features.json')
        
        # 确保目录存在
        os.makedirs(os.path.dirname(master_features_path), exist_ok=True)
        
        try:
            with open(master_features_path, 'w', encoding='utf-8') as f:
                json.dump(self.master_features, f, ensure_ascii=False, indent=2)
            logger.info(f"成功保存名家特征库到: {master_features_path}")
            return True
        except Exception as e:
            logger.error(f"保存名家特征库失败: {e}")
            return False 